module Table where

import Debug.QuickCheck
import Data.List

-------------------------------------------------------------------

data Table a b
  = Join (Table a b) a b (Table a b)
  | Empty
 deriving ( Eq, Show )
 
-------------------------------------------------------------------

lookupT :: Ord a => a -> Table a b -> Maybe b
lookupT x Empty = Nothing
lookupT x (Join left x' y' right)
  | x < x'      = lookupT x left
  | x == x'     = Just y'
  | x > x'      = lookupT x right

-------------------------------------------------------------------

insertT :: Ord a => a -> b -> Table a b -> Table a b
insertT x y Empty = Join Empty x y Empty
insertT x y (Join left x' y' right)
  | x < x'        = Join (insertT x y left) x' y' right
  | x == x'       = Join left x' y right
  | x > x'        = Join left x' y' (insertT x y right)

-------------------------------------------------------------------

-- properties

prop_LookupT :: Int -> Table Int Int -> Bool
prop_LookupT k t = 
  lookupT k t == lookup k (contents t)

prop_InsertT :: Int -> Int -> Table Int Int -> Bool
prop_InsertT k v t = 
  insert (k,v) [(k',v') | (k',v') <- contents t, k' /= k] ==
     contents (insertT k v t)
    
prop_Lookup_Insert :: Int -> Int -> Int -> Table Int Int -> Bool
prop_Lookup_Insert k' k v t =
  lookupT k' (insertT k v t) == 
    if k == k' then Just v else lookupT k' t

-------------------------------------------------------------------

contents :: Table a b -> [(a,b)]
contents Empty          = []
contents (Join l x y r) = contents l ++ [(x,y)] ++ contents r

-------------------------------------------------------------------

-- this is the first, wrong definition
{-
instance (Arbitrary a, Arbitrary b) => Arbitrary (Table a b) where
  arbitrary = sized arbTable
-}

arbTable :: (Arbitrary a, Arbitrary b) => Int -> Gen (Table a b)
arbTable s =
  frequency
  [ (1, return Empty)
  , (s, do x <- arbitrary
           y <- arbitrary
           (l,r) <- two (arbTable (s `div` 2))
           return (Join l x y r))
  ]

-------------------------------------------------------------------

ordered :: Ord a => [a] -> Bool
ordered []       = True
ordered [x]      = True
ordered (x:y:xs) = x <= y && ordered (y:xs)

prop_InvTable :: Table Integer Integer -> Bool
prop_InvTable tab = ordered xs && xs == nub xs
 where
  xs = [ x | (x,y) <- contents tab ]

-------------------------------------------------------------------

instance (Arbitrary a, Ord a, Arbitrary b) => Arbitrary (Table a b) where
  arbitrary = 
    do xys <- arbitrary
       return (table xys)

-- table kvs converts a list of key-value pairs into a Table
-- satisfying the ordering invariant
table :: Ord a => [(a,b)] -> Table a b
table []          = Empty
table ((x,y):xys) = Join (table smaller) x y (table larger)
 where
  smaller = [(x',y') | (x',y') <- xys, x' < x] -- fixed: <= becomes <
  larger  = [(x',y') | (x',y') <- xys, x' > x]

-------------------------------------------------------------------
